package service

import (
	"context"
	"sync"
	"time"

	"gitee.com/linxing_3/sr-core/domain/service/srvctx"
	"gitee.com/linxing_3/sr-core/plugins/log"
	"gitee.com/mkwchecking/sr-admin/apps/knowledge/acl"
	"gitee.com/mkwchecking/sr-admin/apps/knowledge/repo"
	"github.com/ollama/ollama/api"
	"github.com/tmc/langchaingo/chains"
)

func NewRag() (*RAGService, error) {
	// chain, err := repo.NewLangchain()
	// if err != nil {
	// 	return nil, err
	// }

	return &RAGService{
		// retriever:         chain,
		conversationCache: &sync.Map{},
	}, nil
}

type ChatChan struct {
	Msg api.ChatResponse
	Err error
}

type RAGService struct {
	retriever *repo.LangChain
	// docRepo   *repo.KnoDoc
	// groupRepo *repo.KnoDocGroup

	conversationCache *sync.Map
	lock              sync.Mutex
}

func (s *RAGService) CheckRetriever() error {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.retriever == nil {
		chain, err := repo.NewLangchain()
		if err != nil {
			return err
		}
		s.retriever = chain
	}
	return nil
}
func (s *RAGService) Chat(sctx srvctx.IContext, req *acl.RAGChatReq, chatChan chan ChatChan) (string, chan ChatChan, error) {
	// sessinf, err := ctx.GetSessionInfo()
	// if err != nil {
	// 	return "", err
	// }
	if err := s.CheckRetriever(); err != nil {
		return "", nil, err
	}

	ctx := context.Background()

	// key := fmt.Sprintf("%s_%s", sessinf.UserId, req.Namespace)
	var err error
	key := req.Namespace
	conv, ok := s.conversationCache.Load(key)
	if !ok {
		conv, err = s.retriever.GetRetrieverConversation(ctx, repo.DefaultQAPrompt, req.Model, req.Namespace)
		if err != nil {
			return "", nil, err
		}
		s.conversationCache.Store(key, conv)
	}

	userInput := ""
	if len(req.Messages) > 0 {
		userInput = req.Messages[len(req.Messages)-1].Content
	}

	qa := conv.(*chains.RetrievalQA)
	opts := make([]chains.ChainCallOption, 0)
	opts = append(opts, chains.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
		chatChan <- ChatChan{
			Msg: api.ChatResponse{
				Model:     req.Model,
				CreatedAt: time.Now(),
				Message: api.Message{
					Content: string(chunk),
					Role:    "assistant",
				},
				Done: false,
			},
		}
		return nil
	}))

	go func() {
		// defer close(chatChan)
		defer close(chatChan)
		_, err := chains.Run(ctx, qa, userInput, opts...)
		if err != nil {
			chatChan <- ChatChan{
				Err: err,
			}
		}
		log.Debug("run chain done", "err", err)
	}()

	return "", chatChan, nil
}
